#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   ACS.py    
@Contact :   ruaqy@qq.com
@License :   (C)Copyright 2022-~, GPL 2.0

@Modify Time      @Author    @Version    @Description
------------      -------    --------    -----------
2022/7/5 16:50   rqy        1.0         None
"""
import numpy as np
import matplotlib.pyplot as plt


class TSP:
    def __init__(self, city_map, rho):
        self.__CITY_NUM = city_map.shape[0]
        self.__rho = rho

        self.city = city_map

        self.distance_graph = self.get_distance()
        # the initial pheromone is the maximum pheromone
        self.pheromone_graph = np.ones((self.__CITY_NUM, self.__CITY_NUM))

    # calculate the distance among different city
    def get_distance(self):
        distance_graph = []
        for i in range(self.__CITY_NUM):
            city = self.city[i]
            distance = self.city - city
            distance = np.linalg.norm(distance, axis=1)
            distance_graph.append(distance)
        return distance_graph

    def pheromone_update(self, path, distance):
        self.pheromone_graph *= (1 - self.__rho)
        for i in range(self.__CITY_NUM):
            start, end = path[i], path[i + 1]
            self.pheromone_graph[start][end] += self.__rho / distance
            self.pheromone_graph[end][start] += self.__rho / distance


class Ant:
    def __init__(self, alpha, beta, delta, q0, city_num):
        self.__alpha = alpha
        self.__beta = beta
        self.__delta = delta
        self.__q0 = q0

        self.__CITY_NUM = city_num



        self.path = None
        self.distance = None
        self.move_count = None

        self.__clean()

    def __clean(self):
        self.path = [np.random.randint(0, self.__CITY_NUM)]
        self.distance = 0
        self.move_count = 1

    def __choice_next_city(self, distance_graph, pheromone_graph):
        current_city = self.path[-1]
        distance = distance_graph[current_city]
        pheromone = pheromone_graph[current_city]

        # the probability of each city
        city_prob = np.zeros(self.__CITY_NUM)
        for i in range(self.__CITY_NUM):
            if i not in self.path:
                city_prob[i] = np.power(pheromone[i], self.__alpha) * np.power(distance[i], self.__beta)
        city_prob /= np.sum(city_prob)  # uniform
        city_cal_prob = np.cumsum(city_prob)

        q = np.random.random()
        if q > self.__q0:
            prob = np.random.random()
            for i in range(self.__CITY_NUM):
                if city_cal_prob[i] > prob and i not in self.path:
                    break
        else:
            i = np.argmax(city_prob)

        # local updating rule
        local_distance = distance[i]
        pheromone_graph[current_city][i] *= (1 - self.__delta)
        pheromone_graph[current_city][i] += self.__delta / local_distance
        pheromone_graph[i][current_city] *= (1 - self.__delta)
        pheromone_graph[i][current_city] += self.__delta / local_distance
        return i, local_distance

    def update(self, tsp):
        self.__clean()
        while self.move_count < self.__CITY_NUM:
            next_city, distance = self.__choice_next_city(tsp.distance_graph, tsp.pheromone_graph)
            self.path.append(next_city)
            self.distance += distance
            self.move_count += 1
        else:
            distance = tsp.distance_graph[self.path[-1]][self.path[0]]
            self.path.append(self.path[0])
            self.distance += distance



def best_ant_find(ants):
    distance = [ant.distance for ant in ants]
    best_index = np.argmin(distance)
    ant = ants[best_index]
    k_best_path = ant.path
    k_best_distance = ant.distance
    return k_best_path, k_best_distance


def main():
    CITY_NUM, ANTS_NUM, T = 34, 50, 100
    ALPHA, BETA, RHO, DELTA, Q0 = 2, -2, 0.1, 0.0001, 0.9

    city_map = np.array([[116.46, 39.92],
                         [117.2, 39.13],
                         [121.48, 31.22],
                         [106.54, 29.59],
                         [91.11, 29.97],
                         [87.68, 43.77],
                         [106.27, 38.47],
                         [111.65, 40.82],
                         [108.33, 22.84],
                         [126.63, 45.75],
                         [125.35, 43.88],
                         [123.38, 41.8],
                         [114.48, 38.03],
                         [112.53, 37.87],
                         [101.74, 36.56],
                         [117, 36.65],
                         [113.6, 34.76],
                         [118.78, 32.04],
                         [117.27, 31.86],
                         [120.19, 30.26],
                         [119.3, 26.08],
                         [115.89, 28.68],
                         [113, 28.21],
                         [114.31, 30.52],
                         [113.23, 23.16],
                         [121.5, 25.05],
                         [110.35, 20.02],
                         [103.73, 36.03],
                         [108.95, 34.27],
                         [104.06, 30.67],
                         [106.71, 26.57],
                         [102.73, 25.04],
                         [114.1, 22.2],
                         [113.33, 22.13], ])

    tsp = TSP(city_map, RHO)
    ants = [Ant(ALPHA, BETA, DELTA, Q0, CITY_NUM) for i in range(ANTS_NUM)]

    [ant.update(tsp) for ant in ants]
    g_best_path, g_best_distance = best_ant_find(ants)  # global best ant:path and distance

    for k in range(T):
        # search path
        [ant.update(tsp) for ant in ants]

        # update pheromone:permit global best ant
        k_best_path, k_best_distance = best_ant_find(ants)  # find the best ant in the routine

        if k_best_distance < g_best_distance:
            g_best_path = k_best_path
            g_best_distance = k_best_distance

        tsp.pheromone_update(g_best_path, g_best_distance)  # release pheromone

        print('iter', k, k_best_distance)

        plt.clf()
        plt.scatter([i[0] for i in tsp.city], [i[1] for i in tsp.city])
        for i in range(CITY_NUM):
            plt.plot([tsp.city[k_best_path[i % CITY_NUM]][0],
                      tsp.city[k_best_path[i + 1]][0]],
                     [tsp.city[k_best_path[i % CITY_NUM]][1],
                      tsp.city[k_best_path[i + 1]][1]])
        plt.title('iter:{}'.format(k))
        plt.pause(0.01)

    plt.figure()
    plt.scatter([i[0] for i in tsp.city], [i[1] for i in tsp.city])
    for i in range(CITY_NUM):
        plt.plot([tsp.city[k_best_path[i % CITY_NUM]][0],
                  tsp.city[k_best_path[i + 1]][0]],
                 [tsp.city[k_best_path[i % CITY_NUM]][1],
                  tsp.city[k_best_path[i + 1]][1]])
    plt.title('The Best Path:{}'.format(g_best_distance))

    plt.show()


if __name__ == '__main__':
    main()
